/* Copyright (c) Huawei Technologies Co., Ltd. 2023-2023.
 * Description: page attribute check
 * Author: zhangxun <zhangxun38@huawei.com>
 * Maintainer: Xun Zhang
 * Create: 2023-08-15
 */
#include <linux/rbtree.h>
#include <linux/interval_tree.h>
#include <linux/interval_tree_generic.h>
#include <linux/module.h>
#include <linux/slab.h>
#include <linux/types.h>
#include <linux/sysfs.h>
#include <linux/device.h>
#include <linux/securec.h>

#include <linux/hal/pageattr_check.h>

#define INTERVAL_MIN 0
#define INTERVAL_MAX ULLONG_MAX

#define PAGE_ATTR_PRINT_MAX     500
#define PAGE_ATTR_PRINT_DEFAULT 5
#define DEFAULT_ENABLE 0

struct pageattr_module_data {
	struct kobject *kobj;
	struct rb_root_cached root;
	int enable;
};

struct page_attr_t {
	struct list_head list;
	struct rb_node rb;
	unsigned long long start;
	unsigned long long size;
	unsigned long long attr;
	unsigned long long mask;
	unsigned long long __subtree_last;
};

#define START(node) ((node)->start)
#define LAST(node)  ((node)->start + (node)->size - 1)
INTERVAL_TREE_DEFINE(struct page_attr_t, rb, unsigned long long, __subtree_last,
		     START, LAST, static inline, pageattr);

#define for_all_interval_tree_node_range(tree, node, start, size) \
		for (node = pageattr_iter_first(tree, start, size); node; \
		node = pageattr_iter_next(node, start, size))

#define for_all_interval_tree_node(tree, node) \
		for_all_interval_tree_node_range(tree, node, INTERVAL_MIN, INTERVAL_MAX)

static struct pageattr_module_data *g_pageattr;
static atomic_t g_guard = ATOMIC_INIT(0);
static atomic_t g_printnum = ATOMIC_INIT(PAGE_ATTR_PRINT_DEFAULT);

struct pageattr_module_data *get_pageattr_data(void)
{
	return g_pageattr;
}

void pageattr_set_printnum(int num)
{
	atomic_set(&g_printnum, num);
}

int pageattr_get_printnum(void)
{
	return atomic_read(&g_printnum);
}

struct page_attr_t *page_attr_get(unsigned long long start, unsigned long long size)
{
	if (unlikely(!g_pageattr))
		return NULL;

	return pageattr_iter_first(&g_pageattr->root, start, size);
}

int pageattr_insert_range(unsigned long long start, unsigned long long size,
	unsigned long long attr, unsigned long long mask)
{
	struct page_attr_t *new;

	new = kzalloc(sizeof(struct page_attr_t), GFP_KERNEL);
	if (!new)
		return -ENOMEM;

	new->start = start;
	new->size  = size;
	new->attr  = attr;
	new->mask  = mask;
	atomic_set(&g_guard, 1);
	pageattr_insert(new, &g_pageattr->root);
	atomic_set(&g_guard, 0);
	return 0;
}

/*
 * There are total 4 situations to handle:
 * |  (new)  |   (new)  (exist)  | (exist) | start < exist->start
 * | (exist) |  (exist)  (new)   |  (new)  | last > exist->last
 * | (exist) |  (exist)  (new)   | (exist) |
 * |  (new)  |   (new)  (exist)  |  (new)  | start < exist->start, last > exist->last
 */
static int pageattr_register_range_merge(unsigned long long start, unsigned long long size,
	unsigned long long attr, unsigned long long mask, struct page_attr_t *exist)
{
	pr_err("range coincides, not support!\n");
	return -EINVAL;
}

int pageattr_register_range(unsigned long long start, unsigned long long size,
	unsigned long long attr, unsigned long long mask)
{
	struct page_attr_t *exist;

	if (unlikely(!g_pageattr))
		return -ENOMEM;

	if (size == 0) {
		pr_err("size should not be 0!\n");
		return -EINVAL;
	}

	if ((start + size) < start) {
		pr_err("input range size overflow!\n");
		return -EOVERFLOW;
	}

	exist = page_attr_get(start, size);
	if (exist)
		return pageattr_register_range_merge(start, size, attr, mask, exist);

	return pageattr_insert_range(start, size, attr, mask);
}
EXPORT_SYMBOL(pageattr_register_range);

void pageattr_clear_all(void)
{
	struct page_attr_t *node, *next;
	LIST_HEAD(to_remove);

	if (unlikely(!g_pageattr))
		return;

	atomic_set(&g_guard, 1);
	for_all_interval_tree_node(&g_pageattr->root, node) {
		pageattr_remove(node, &g_pageattr->root);
		list_add(&node->list, &to_remove);
	}

	/* Nothing to remove */
	if (list_empty(&to_remove)) {
		atomic_set(&g_guard, 0);
		return;
	}

	list_for_each_entry_safe(node, next, &to_remove, list) {
		list_del_init(&node->list);
		kfree(node);
	}
	atomic_set(&g_guard, 0);
}
EXPORT_SYMBOL(pageattr_clear_all);

/* pxx_val is the value from pte/pmd/pud */
int pageattr_check(unsigned long long pxx_val, unsigned long long size)
{
	struct page_attr_t *data;
	unsigned long long start = pxx_val & GENMASK_ULL(52, 12);

	if (unlikely(!g_pageattr))
		return 0;

	if (!g_pageattr->enable)
		return 0;

	if (unlikely(atomic_read(&g_guard)))
		return 0;

	data = page_attr_get(start, size);
	if (!data)
		return 0;

	if (likely((pxx_val & data->mask) == data->attr))
		return 0;

	if (atomic_read(&g_printnum) <= 0)
		return -1;

	pr_warn("WARNING: [%llx, %llx] expect:%llx->%llx pxx_val:%llx\n",
			start, size, data->attr, pxx_val & data->mask, pxx_val);
	atomic_dec(&g_printnum);
	return -1;
}

static ssize_t printnum_show(struct device *dev,
							struct device_attribute *attr, char *buf)
{
	struct pageattr_module_data *data;

	data = get_pageattr_data();
	if (!data)
		return 0;

	return sprintf_s(buf, PAGE_SIZE, "%d\n", pageattr_get_printnum());
}

static ssize_t printnum_store(struct device *dev,
							struct device_attribute *attr,
							const char *buf, size_t buf_size)
{
	int num;
	struct pageattr_module_data *data;

	data = get_pageattr_data();
	if (!data)
		return -ENOMEM;

	/* need to match 1 argument */
	if (sscanf_s(buf, "%d", &num) != 1)
		return -EINVAL;

	if (num > PAGE_ATTR_PRINT_MAX || num < 0) {
		pr_err("pageattr input invalid: num should be 0-%d\n", PAGE_ATTR_PRINT_MAX);
		return -EINVAL;
	}

	pageattr_set_printnum(num);
	return buf_size;
}
DEVICE_ATTR_ADMIN_RW(printnum);

static ssize_t enable_show(struct device *dev,
							struct device_attribute *attr, char *buf)
{
	struct pageattr_module_data *data;

	data = get_pageattr_data();
	if (!data)
		return 0;

	return sprintf_s(buf, PAGE_SIZE, "%d\n", data->enable);
}

static ssize_t enable_store(struct device *dev,
							struct device_attribute *attr,
							const char *buf, size_t buf_size)
{
	int enable;
	struct pageattr_module_data *data;

	data = get_pageattr_data();
	if (!data)
		return -ENOMEM;

	/* need to match 1 argument */
	if (sscanf_s(buf, "%d", &enable) != 1)
		return -EINVAL;

	if (enable != 0 && enable != 1) {
		pr_err("pageattr input invalid: enable should be 1 or 0\n");
		return -EINVAL;
	}

	data->enable = enable;
	return buf_size;
}
DEVICE_ATTR_ADMIN_RW(enable);

static struct attribute *pageattr_attributes[] = {
	&dev_attr_enable.attr,
	&dev_attr_printnum.attr,
	NULL
};

static const struct attribute_group pageattr_attr_group = {
	.attrs = pageattr_attributes,
};

void page_attr_sys_exit(void)
{
	struct pageattr_module_data *data;

	data = get_pageattr_data();
	if (!data)
		return;

	sysfs_remove_group(data->kobj, &pageattr_attr_group);
	kobject_put(data->kobj);
}

int __init page_attr_sys_init(void)
{
	int ret;
	struct pageattr_module_data *data;

	data = get_pageattr_data();
	if (!data)
		return -ENOMEM;

	data->kobj = kobject_create_and_add("pageattr", kernel_kobj);

	if (!data->kobj)
		return -ENOMEM;

	ret = sysfs_create_group(data->kobj, &pageattr_attr_group);
	if (ret < 0) {
		pr_err("Failed to create attribute sysfs\n");
		kobject_put(data->kobj);
		return ret;
	}
	return 0;
}

static void __exit page_attr_exit(void)
{
	page_attr_sys_exit();
	pageattr_clear_all();
	kfree(g_pageattr);
	g_pageattr = NULL;
}

int __init page_attr_init(void)
{
	return page_attr_sys_init();
}

int __init page_attr_early_init(void)
{
	g_pageattr = kzalloc(sizeof(struct pageattr_module_data), GFP_KERNEL);

	if (!g_pageattr)
		return -ENOMEM;

	g_pageattr->root     = RB_ROOT_CACHED;
	g_pageattr->enable   = DEFAULT_ENABLE;

	pr_info("%s succeed\n", __func__);
	return 0;
}

early_initcall(page_attr_early_init);
module_init(page_attr_init);
module_exit(page_attr_exit);
